{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 测试相同形状模式"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.script import relax as R\n",
    "from tvm.script import tir as T\n",
    "from tvm import relax as rx\n",
    "from tvm import relay, tir\n",
    "from tvm.relax.analysis import get_var2val\n",
    "import tvm.testing\n",
    "from tvm.relax.dpl import *\n",
    "\n",
    "same_shape_func_type = tvm.testing.parameter(\n",
    "    \"same_static_shape\",\n",
    "    \"same_dynamic_shape\",\n",
    "    \"different_static_shape\",\n",
    "    \"different_dynamic_shape\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_same_shape_pattern(same_shape_func_type):\n",
    "    if same_shape_func_type == \"same_static_shape\":\n",
    "\n",
    "        @R.function(private=True)\n",
    "        def func(\n",
    "            a: R.Tensor((1024, 128), \"float32\"),\n",
    "            b: R.Tensor((1024, 128), \"float32\"),\n",
    "        ) -> R.Tensor:\n",
    "            with R.dataflow():\n",
    "                c = R.multiply(a, R.const(2.0))\n",
    "                d = R.add(b, c)\n",
    "                out = d\n",
    "                R.output(out)\n",
    "            return out\n",
    "\n",
    "    elif same_shape_func_type == \"same_dynamic_shape\":\n",
    "\n",
    "        @R.function(private=True)\n",
    "        def func(\n",
    "            a: R.Tensor((\"n\", 128), \"float32\"),\n",
    "            b: R.Tensor((\"n\", 128), \"float32\"),\n",
    "        ) -> R.Tensor:\n",
    "            with R.dataflow():\n",
    "                c = R.multiply(a, R.const(2.0))\n",
    "                d = R.add(b, c)\n",
    "                out = d\n",
    "                R.output(out)\n",
    "            return out\n",
    "\n",
    "    elif same_shape_func_type == \"different_static_shape\":\n",
    "\n",
    "        @R.function(private=True)\n",
    "        def func(\n",
    "            a: R.Tensor((1024, 128), \"float32\"),\n",
    "            b: R.Tensor((1, 128), \"float32\"),\n",
    "        ) -> R.Tensor:\n",
    "            with R.dataflow():\n",
    "                c = R.multiply(a, R.const(2.0))\n",
    "                d = R.add(b, c)\n",
    "                out = d\n",
    "                R.output(out)\n",
    "            return out\n",
    "\n",
    "    elif same_shape_func_type == \"different_dynamic_shape\":\n",
    "\n",
    "        @R.function(private=True)\n",
    "        def func(\n",
    "            a: R.Tensor((\"n\", 128), \"float32\"),\n",
    "            b: R.Tensor((\"m\", 128), \"float32\"),\n",
    "        ) -> R.Tensor:\n",
    "            with R.dataflow():\n",
    "                c = R.multiply(a, R.const(2.0))\n",
    "                d = R.add(b, c)\n",
    "                out = d\n",
    "                R.output(out)\n",
    "            return out\n",
    "\n",
    "    else:\n",
    "        raise ValueError(f\"Unknown value of same_shape_func_type={same_shape_func_type}\")\n",
    "\n",
    "    with PatternContext() as ctx:\n",
    "        pat_lhs = wildcard()\n",
    "        pat_rhs = wildcard()\n",
    "        pat_sum = is_op(\"relax.add\")(pat_lhs, pat_rhs)\n",
    "        pat_lhs.same_shape_as(pat_rhs)\n",
    "\n",
    "    block = func.body.blocks[0]\n",
    "    match = ctx.match_dfb(block)\n",
    "\n",
    "    if \"same\" in same_shape_func_type:\n",
    "        assert match\n",
    "    else:\n",
    "        assert match is None\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ai",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
